﻿#pragma once

#include "vec_n.h"
#include <algorithm>
#include <cmath>

namespace ifire {

template <int M, int N> class MatMN {
private:
  VecN<N> rows[M]; // 静态数组存储行

public:
  // 构造函数
  MatMN() {
    static_assert(M > 0 && N > 0, "Matrix dimensions must be positive");
  }

  // 禁用拷贝构造/赋值（如果需要深拷贝，可以自行实现）
  MatMN(const MatMN&) = delete;
  MatMN& operator=(const MatMN&) = delete;

  // 允许移动
  MatMN(MatMN&&) = default;
  MatMN& operator=(MatMN&&) = default;

  // 矩阵置零
  void Zero() {
    for (int i = 0; i < M; ++i) {
      rows[i].Zero();
    }
  }

  // 矩阵转置
  template <int P = N, int Q = M> MatMN<P, Q> Transpose() const {
    static_assert(P == N && Q == M, "Transpose dimensions must match");
    MatMN<P, Q> result;
    for (int i = 0; i < M; ++i) {
      for (int j = 0; j < N; ++j) {
        result[j][i] = rows[i][j];
      }
    }
    return result;
  }

  // 矩阵-向量乘法
  template <int P> VecN<M> operator*(const VecN<P>& v) const {
    static_assert(N == P, "Vector dimension must match matrix columns");
    VecN<M> result;
    for (int i = 0; i < M; ++i) {
      result[i] = rows[i].Dot(v);
    }
    return result;
  }

  // 矩阵-矩阵乘法
  template <int P> MatMN<M, P> operator*(const MatMN<N, P>& m) const {
    MatMN<M, P> result;
    for (int i = 0; i < M; ++i) {
      for (int j = 0; j < P; ++j) {
        float sum = 0.0f;
        for (int k = 0; k < N; ++k) {
          sum += rows[i][k] * m[k][j];
        }
        result[i][j] = sum;
      }
    }
    return result;
  }

  // 高斯-赛德尔求解
  static VecN<M> SolveGaussSeidel(const MatMN<M, N>& A, const VecN<N>& b) {
    static_assert(M == N, "Matrix must be square for Gauss-Seidel");
    VecN<N> X;
    X.Zero();

    const int MAX_ITERATIONS = 50;
    const float TOLERANCE = 1e-5f;

    for (int iter = 0; iter < MAX_ITERATIONS; ++iter) {
      float max_diff = 0.0f;
      for (int i = 0; i < N; ++i) {
        if (fabs(A.rows[i][i]) < FLT_EPSILON) continue;

        float sigma = 0.0f;
        for (int j = 0; j < N; ++j) {
          if (j != i) sigma += A.rows[i][j] * X[j];
        }

        float new_val = (b[i] - sigma) / A.rows[i][i];
        max_diff = (std::max)(max_diff, fabs(new_val - X[i]));
        X[i] = new_val;
      }

      if (max_diff < TOLERANCE) break;
    }
    return X;
  }

  // 访问行
  VecN<N>& operator[](int index) { return rows[index]; }
  const VecN<N>& operator[](int index) const { return rows[index]; }
};

} // namespace ifire